from flask import Flask, jsonify, request
import requests
from flask_swagger_ui import get_swaggerui_blueprint  # 引入 Swagger UI

app = Flask(__name__)
app.config['JSON_SORT_KEYS'] = False  # 禁用键排序

# 定义外部接口URL和请求头
UPLOAD_URL = "http://10.252.202.55/v1/files/upload"
CHAT_MESSAGES_URL = "http://10.252.202.55/v1/chat-messages"
WORKFLOWS_RUN_URL = "http://10.252.202.55/v1/workflows/run"

UPLOAD_HEADERS = {
    "Authorization": "Bearer app-SEUTdEhQzosL417F6qcJoGH9"
}

LAW_MESSAGES_HEADERS = {
    "Authorization": "Bearer app-KGDFk5qG1lsSqVDFCx9SQvhZ",
    "Content-Type": "application/json"
}

ZHIDU_MESSAGES_HEADERS = {
    "Authorization": "Bearer app-0GPbIHj6zzkBjt8m4Jl7cMpl",
    "Content-Type": "application/json"
}

EXTRACT_REVIEW_HEADERS = {
    "Authorization": "Bearer app-SEUTdEhQzosL417F6qcJoGH9",
    "Content-Type": "application/json"
}

REVIEW_HEADERS = {
    "Authorization": "Bearer app-XaTng8MSYczqcIHtOvjcWO8y",
    "Content-Type": "application/json"
}

GENERATE_HEADERS = {
    "Authorization": "Bearer app-rX9foBHM5Lir1uwNEz2t5nIi",
    "Content-Type": "application/json"
}

# 上传文件
@app.route('/v1/upload', methods=['POST'])
def upload():
    # 定义外部接口URL和请求头
    url = UPLOAD_URL
    headers = UPLOAD_HEADERS
    
    try:
        # 获取文件和用户标识
        file = request.files['file']
        user = request.form['user']
        
        if not file or not user:
            return jsonify({"error": "缺少文件或用户标识"}), 400
        
        # 创建 multipart/form-data 请求
        files = {'file': (file.filename, file, file.content_type)}  # 确保文件对象正确
        data = {'user': user}
        
        # 调用外部接口
        response = requests.post(url, headers=headers, files=files, data=data)
        
        if response.status_code == 201:
            return jsonify(response.json())
        else:
            return jsonify({"error": "文件上传失败", "details": response.json()}), response.status_code
    except Exception as e:
        return jsonify({"error": "上传过程中发生错误", "details": str(e)}), 500




# 法律助手对话
@app.route('/v1/law-assistant', methods=['POST'])
def chat_law():
    # 定义外部接口URL和请求头
    url = CHAT_MESSAGES_URL
    headers = LAW_MESSAGES_HEADERS
    try:
        payload = request.json
        if not payload:
            return jsonify({"error": "缺少参数"}), 400
    except Exception:
        return jsonify({"error": "无效的请求体，请用json请求"}), 400

    # 调用外部接口
    response = requests.post(url, json=payload, headers=headers, stream=True)
    if response.status_code != 200:
        return jsonify({"error": "外部接口调用失败"}), response.status_code

    # 返回流式响应
    def generate():
        for chunk in response.iter_content(chunk_size=8192, decode_unicode=True):
            yield chunk

    return app.response_class(generate(), mimetype='text/event-stream')

# 制度助手对话
@app.route('/v1/institution-assistant', methods=['POST'])
def chat_institution():
    # 定义外部接口URL和请求头
    url = CHAT_MESSAGES_URL
    headers = ZHIDU_MESSAGES_HEADERS
    try:
        payload = request.json
        if not payload:
            return jsonify({"error": "缺少参数"}), 400
    except Exception:
        return jsonify({"error": "无效的请求体，请用json请求"}), 400

    # 调用外部接口
    response = requests.post(url, json=payload, headers=headers, stream=True)
    if response.status_code != 200:
        return jsonify({"error": "外部接口调用失败"}), response.status_code

    # 返回流式响应
    def generate():
        for chunk in response.iter_content(chunk_size=8192, decode_unicode=True):
            yield chunk

    return app.response_class(generate(), mimetype='text/event-stream')

# 合同抽取
@app.route('/v1/extract', methods=['POST'])
def extract():

    # 定义外部接口URL和请求头
    url = WORKFLOWS_RUN_URL
    headers = EXTRACT_REVIEW_HEADERS
    try:
        payload = request.json
        if not payload:
            return jsonify({"error": "缺少参数"}), 400
    except Exception:
        return jsonify({"error": "无效的请求体，请用json请求"}), 400

    # 调用外部接口
    response = requests.post(url, json=payload, headers=headers)
    try:
        # 提取返回结果中的 data.outputs.text 字段
        result = response.json()
        text_data = result.get("data", {}).get("outputs", {}).get("text", [])
        if text_data:
            # 直接提取并返回 JSON 格式的数据
            contract_info = {
                "contract_name": text_data[0].get("合同名称"),
                "contract_subject": text_data[0].get("合同主体"),
                "partner_company": text_data[0].get("合作单位"),
                "contract_amount": text_data[0].get("合同金额"),
                "start_date": text_data[0].get("开始时间"),
                "end_date": text_data[0].get("结束时间"),
                "contract_target": text_data[0].get("合同标的"),
                "main_content": text_data[0].get("合同主要内容")
            }
            return jsonify(contract_info)
        else:
            return jsonify({"error": "文本生成规则不对，无法提取，请联系后端"}), 400
    except ValueError:
        return jsonify({"error": "返回内容不是有效的JSON格式，请联系后端"}), 400


# 合同审查
@app.route('/v1/review', methods=['POST'])
def review():
    # 定义外部接口URL和请求头
    url = WORKFLOWS_RUN_URL
    headers = REVIEW_HEADERS
    try:
        payload = request.json
        if not payload:
            return jsonify({"error": "缺少参数"}), 400
    except Exception:
        return jsonify({"error": "无效的请求体，请用json请求"}), 400

    # 调用外部接口
    response = requests.post(url, json=payload, headers=headers)
    try:
        # 提取返回结果中的 data.outputs.text 字段
        result = response.json()
        text_data = result.get("data", {}).get("outputs", {}).get("text", [])
        if text_data:
            # 根据返回结果生成特定格式的对象
            review_results = []

            # 假设 text_data 中的每个元素包含 category 和 result 字段
            for item in text_data:
                category = item.get("category")
                result_details = item.get("result", [])

                category_result = {
                    "category": category,
                    "result": []
                }

                for detail in result_details:
                    question_info = {
                        "title": detail.get("title"),
                        "original": detail.get("original"),
                        "according": detail.get("according"),
                        "advice": detail.get("advice")
                    }
                    category_result["result"].append(question_info)

                review_results.append(category_result)

            return jsonify(review_results)
        else:
            return jsonify({"error": "文本生成规则不对，无法提取，请联系后端"}), 400
    except ValueError:
        return jsonify({"error": "返回内容不是有效的JSON格式，请联系后端"}), 400

# 智能报告生成
@app.route('/v1/generate', methods=['POST'])
def generate():
    # 定义外部接口URL和请求头
    url = WORKFLOWS_RUN_URL
    headers = GENERATE_HEADERS
    try:
        payload = request.json
        if not payload:
            return jsonify({"error": "缺少参数"}), 400
    except Exception:
        return jsonify({"error": "无效的请求体，请用json请求"}), 400

    # 调用外部接口
    response = requests.post(url, json=payload, headers=headers)
    try:
        # 提取返回结果中的 data.outputs.text 字段
        result = response.json()
        text_data = result.get("data", {}).get("outputs", {}).get("text", [])
        if text_data:
            return jsonify({"text": text_data})
        else:
            return jsonify({"error": "文本生成规则不对，无法提取，请联系后端"}), 400
    except ValueError:
        return jsonify({"error": "返回内容不是有效的JSON格式，请联系后端"}), 400



# 配置 Swagger UI
SWAGGER_URL = '/swagger'
API_URL = '/static/swagger.json'

swaggerui_blueprint = get_swaggerui_blueprint(
    SWAGGER_URL,
    API_URL,
    config={
        'app_name': "大模型后端接口"
    }
)
app.register_blueprint(swaggerui_blueprint, url_prefix=SWAGGER_URL)


if __name__ == '__main__':
    app.run(host='10.1.136.49', port=5000)